from mootdx.affair import Affair
from mootdx.quotes import Quotes
import pandas
import baostock as bs
from mootdx import consts
import struct

# 远程文件列表
#files = Affair.files()
#print(files)

import pymongo
#db_target='mongodb://admin:123.123.@localhost:27017/'

db_target='mongodb://192.168.10.5:27017/'
'''
company_info
contract_future:_id:exchange+symbol,exchange:str,symbol:str,name:str,company_id:int,precision:int32,listed:bool
contract_stock_cn
contract_preclose :_id,int32
contract_etf
'''

def getExchangeID(ex):
  if ex=="SZ":
    return 0
  if ex=="SH":
    return 1


def getMongo():
  return pymongo.MongoClient(db_target)

def getDB(name:str):
  return getMongo()[name]

def get_grouped_symbols():
    client = Quotes.factory(market='std')
    
    ans={"stock":[],"index":[],"bond":[],"fund":[]}
    markets=[consts.MARKET_SH,consts.MARKET_SZ]

    for m in markets:
      symbol = client.stocks(market=m)
      for row in symbol.iterrows():

          unused=[]
          record={'name':row[1]['name'],'slotcnt':row[1]['volunit'],'decimal_point':row[1]['decimal_point']}
          if m ==consts.MARKET_SH:
            record["exchange"]="SH"
          else:
            record["exchange"]="SZ"

    if len(unused)>0:
        print("{} 个symbol无法识别类型:{}".format(len(unused),unused))
    return ans


'''
创建
'''
def update_types():
  db=getDB("contract")
  db["types"].drop()
  markets=[{"_id":0,"code":"STK","name":"股票"},
  {"_id":1,"code":"IND","name":"指数"},
  {"_id":2,"code":"ETF","name":"基金"},
  {"_id":3,"code":"CBD","name":"可转债"}, 

  ]
  db["types"].insert_many(markets)
'''
只根据市场和symbol，确定合约类型
'''
def getType(market:str,symbol:str):
  sh={'STK':['60','68'],'IND':['000'],'CBD':['11']}
  sz={'STK':['00','30'],'IND':['39'],'CBD':['12']}
  etf=['512','515','159','516']

  if market=="SH":
    for key in sh.keys():
      if symbol[0:2] in sh[key] or symbol[0:3] in sh[key]:
        return key
  if market=="SZ":
    for key in sz.keys():
      if symbol[0:2] in sz[key] or symbol[0:3] in sz[key]:
        return key
  if symbol[0:2] in etf or symbol[0:3] in etf:
      return "ETF"
  return "?"

#预处理扩展市场
def caculate_ext_row(df):
  df.drop("desc",axis=1,inplace=True)


#计算集合中每只股票的id，以及属于 类型（股票？债券？ETF？）
def caculate_row(df):
  df["_id"]=0
  df["type_code"]=""
  print("start prepare row")
  for rowi in range(df.shape[0]):
    '''
    if False and df.loc[rowi,"symbol"]=="999999" or df.loc[rowi,"symbol"]=="999998":
      df.drop(rowi,inplace=True)
      continue
    '''
    #计算64bit id

    rawb=bytearray([getExchangeID(df.loc[rowi,"exchange"])])+bytearray(df.loc[rowi,"symbol"].encode('utf-8'))
    if len(rawb)<8:
      arr=[0 for i in range(8-len(rawb))]
      rawb=bytearray(arr)+rawb
    b=struct.unpack(">Q", rawb)
    df.loc[rowi,"_id"]=b
    #完成id计算
    #添加type code
    df.loc[rowi,"type_code"]=getType(df.loc[rowi,"exchange"],df.loc[rowi,"symbol"])
  print("%d rows prepared.",df.shape[0])
    
'''
仅针对常规沪深市场。
传入df，tpe指明了合约类型，本函数根据type，将合约存储到Contract数据库的不同的集合。
'''

def insertByType(df):
  dbs=['STK','ETF','CBD','IND']
  db=getDB("contract")
  for name in dbs:
    print("start drop and insert "+name)
    db[name].drop()
    records=df.loc[df["type_code"]==name].to_dict(orient='records')
    for i in range(len(records)):
      if db[name].find_one({"_id":records[i]["_id"]}):
        continue
      del records[i]["type_code"]

      db[name].insert_one(records[i])

  records=df.loc[df["type_code"]=="?"].to_dict(orient='records')
  for i in range(len(records)):
    if db["unknown"].find_one({"_id":records[i]["_id"]}):
      continue
    del records[i]["type_code"]
    db["unknown"].insert_one(records[i])

#爬取数据并存储的入口函数
def pull_save2mongodb(client):

  symbol = client.stocks(market=consts.MARKET_SH)
  symbol.rename(columns={'code':'symbol','volunit':'trade_min','decimal_point':'precision'},inplace=True)
  symbol["exchange"]="SH"
  symbol.drop("pre_close",axis=1,inplace=True)
  
  symbol_sz = client.stocks(market=consts.MARKET_SZ)
  symbol_sz.rename(columns={'code':'symbol','volunit':'trade_min','decimal_point':'precision'},inplace=True)
  symbol_sz["exchange"]="SZ"
  symbol_sz.drop("pre_close",axis=1,inplace=True)

  symbol=pandas.concat([symbol, symbol_sz], ignore_index=True)
  caculate_row(symbol)
  insertByType(symbol)



def baostock_login():
  # 登陆系统
  lg = bs.login()
# 显示登陆返回信息
  print('login respond error_code:'+lg.error_code)
  print('login respond  error_msg:'+lg.error_msg)

def add_log(m,c=None):
  db=getDB("logger")
  if c:
    db[c].insert_one(m)
  else:
    db["default"].insert_one(m)

def add_listing_date():
  baostock_login()
  db=getDB("contract")
  cur=db["STK"].find({},{"_id":1,"exchange":1,"symbol":1,"listing_date":1})
  wait_save=[{}]
  for row in cur:
    if "listing_date" in row and type(row["listing_date"])==type(0) and row["listing_date"]>0:
      continue
    code= row["exchange"].lower()+"."+row["symbol"]
    rs = bs.query_stock_basic(code=code)
    if rs.error_code!="0":
      add_log({"msg":"baostock error","code":code,"error_code":rs.error_code,"error_msg":rs.error_msg})
      continue
    try:
      d=rs.get_row_data()[2].split("-")
      date=int(d[0])*10000+int(d[1])*100+int(d[2])
      db["STK"].update_one({"_id":row["_id"]},{"$set":{"listing_date":date}})
    except BaseException:
      add_log({"msg":"baostock data error","code":code,"row_data":rs.get_row_data(),"error_code":rs.error_code,"error_msg":rs.error_msg})
      continue
#保存固定的市场。并读取tdx市场数据，保存到contract数据库
def update_markets():
  db=getDB("contract")
  db["markets"].drop()
  markets=[{"_id":0,"code":"SH","name":"上海证券交易所"},
  {"_id":1,"code":"SZ","name":"深圳证券交易所"},
  {"_id":2,"code":"HKE","name":"香港证券交易所"},
  {"_id":10,"code":"CFFEX","name":"中国金融期货交易所"},  #个股期权，ETF期权
  {"_id":11,"code":"SHFE","name":"上海期货交易所"},
  {"_id":12,"code":"DCE","name":"大连商品交易所"},
  {"_id":13,"code":"CZCE","name":"郑州商品交易所"},
  {"_id":14,"code":"INE","name":"上海能源交易所"},
  ]
  db["markets"].insert_many(markets)
  
  client = Quotes.factory(market='ext')
  mkt=client.markets()
  if len(mkt)>0:
    for i in range(mkt.shape[0]):
      mkt.loc[i,"market"]=int(mkt.loc[i,"market"])
    db["tdx_markets"].drop()
    mkt.rename(columns={'category':'type_code',"market":"_id"},inplace=True)
    db["tdx_markets"].insert_many(mkt.to_dict(orient="records"))
  else:
    print("get 0 market info from tdx.")


def save_raw_ext_contracts():
  client = Quotes.factory(market='ext')
  print("markets:",client.markets())
  db=getDB("contract")
  db["ext_raw_contracts"].drop()#删除此集合
  db["ext_raw_contracts"].insert_many(client.instruments().to_dict(orient="records"))

if __name__ == "__main__":
  '''
  from tqsdk import TqApi, TqAuth
  api = TqApi(auth=TqAuth("gogo001", "gogo001"))
  klines=api.get_kline_serial("SHFE.cu2201", 86400)
  print (klines)
  while True:
    api.wait_update()
    print(klines.iloc[-1])
  '''
  update_types()
  #update_markets()#保存市场数据
  #client = Quotes.factory(market='std', multithread=False, heartbeat=False)
  #pull_save2mongodb(client)#保存全市场股票数据
  #save_raw_ext_contracts()
  #client = Quotes.factory(market='ext', multithread=False, heartbeat=False)
  #print(client.markets())
  #print(client.bars(frequency=consts.KLINE_DAILY, market=47, symbol="47#IF2109", start=0, offset=100))
  #print(client.bars(frequency=consts.KLINE_DAILY, market=1, symbol="000001", start=0, offset=100))
  